{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/test-env/bin/python\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "print(sys.executable)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [],
   "source": [
    "import unittest\n",
    "import pymysql\n",
    "from typing import List, Dict, Tuple\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SQLRewriter:\n",
    "    def __init__(self, db_config):\n",
    "        self.db_config = db_config\n",
    "    \n",
    "    def _get_connection(self):\n",
    "        \"\"\"Helper function to create a connection to the database\"\"\"\n",
    "        return pymysql.connect(**self.db_config)\n",
    "\n",
    "    def get_table_privileges(self, tables: List[str], user_id: str) -> Tuple[Dict[str, Dict], str]:\n",
    "        \"\"\"\n",
    "        从权限表中获取用户对指定表的权限\n",
    "\n",
    "        Args:\n",
    "            tables: 表名列表\n",
    "            user_id: 用户ID\n",
    "\n",
    "        Returns:\n",
    "            Tuple[Dict[str, Dict], str]: \n",
    "                - 权限上下文字典，格式如：\n",
    "                  {\n",
    "                    \"table_name\": {\n",
    "                        \"allowed_columns\": [\"col1\", \"col2\"]\n",
    "                    }\n",
    "                  }\n",
    "                - 错误信息\n",
    "        \"\"\"\n",
    "        try:\n",
    "            conn = self._get_connection()\n",
    "            cursor = conn.cursor(pymysql.cursors.DictCursor)\n",
    "            \n",
    "            # 构建表名的占位符\n",
    "            table_placeholders = ','.join(['%s'] * len(tables))\n",
    "            \n",
    "            # 查询用户对这些表的权限\n",
    "            query = f\"\"\"\n",
    "            SELECT \n",
    "                table_name,\n",
    "                allowed_columns\n",
    "            FROM test_db.user_permissions\n",
    "            WHERE user_id = %s\n",
    "                AND table_name IN ({table_placeholders})\n",
    "            \"\"\"\n",
    "            \n",
    "            # 执行查询\n",
    "            cursor.execute(query, [user_id] + tables)\n",
    "            rows = cursor.fetchall()\n",
    "            \n",
    "            # 构建权限上下文\n",
    "            privileges = {}\n",
    "            for table in tables:\n",
    "                privileges[table] = {\n",
    "                    \"allowed_columns\": [\"*\"],  # 默认允许所有列\n",
    "                }\n",
    "            \n",
    "            # 更新查询到的权限\n",
    "            for row in rows:\n",
    "                table_name = row['table_name']\n",
    "                if table_name in privileges:\n",
    "                    # 处理列权限\n",
    "                    if row['allowed_columns']:\n",
    "                        privileges[table_name][\"allowed_columns\"] = row['allowed_columns'].split(',')\n",
    "            \n",
    "            cursor.close()\n",
    "            conn.close()\n",
    "            \n",
    "            return privileges, \"\"\n",
    "        \n",
    "        except Exception as e:\n",
    "            error_msg = f\"获取表权限时发生错误: {str(e)}\"\n",
    "            print(error_msg)\n",
    "            return {}, error_msg\n",
    "\n",
    "    def rewrite_sql_with_privileges(self, sql: str, user_id: str) -> Tuple[Optional[str], str]:\n",
    "        \"\"\"\n",
    "        根据用户权限重写 SQL 语句\n",
    "        Args:\n",
    "            sql: 原始 SQL 语句\n",
    "            user_id: 用户ID\n",
    "        Returns:\n",
    "            Tuple[Optional[str], str]: (重写后的 SQL, 错误信息)\n",
    "        \"\"\"\n",
    "        try:\n",
    "            # 1. 提取 SQL 中的表名\n",
    "            tables = self.extract_tables_from_sql(sql)\n",
    "            if not tables:\n",
    "                return None, \"无法解析 SQL 中的表名\"            \n",
    "            table_name = tables[0]  # 假设 SQL 只涉及一个表\n",
    "            \n",
    "            # 2. 获取权限信息\n",
    "            privileges, error = self.get_table_privileges(tables, user_id)\n",
    "            if error:\n",
    "                return None, error\n",
    "            \n",
    "            # 3. 获取允许访问的列\n",
    "            allowed_columns = privileges.get(table_name, {}).get('allowed_columns', [])\n",
    "            if not allowed_columns or allowed_columns == [\"*\"]:\n",
    "                return None, f\"权限不足，无法访问表 {table_name}\"\n",
    "            \n",
    "            # 4. 重写 SQL 查询，只保留允许的列\n",
    "            rewritten_sql = sql.replace(\"*\", \", \".join([col.strip() for col in allowed_columns]))\n",
    "            \n",
    "            # 5. 返回重写后的 SQL\n",
    "            return rewritten_sql, \"\"\n",
    "        except Exception as e:\n",
    "            return None, f\"SQL 重写过程出错: {str(e)}\"\n",
    "\n",
    "    def extract_tables_from_sql(self, sql: str) -> list:\n",
    "        \"\"\"简单的从 SQL 中提取表名（假设是 'FROM table_name' 格式）\"\"\"\n",
    "        sql = sql.lower()\n",
    "        if \"from\" in sql:\n",
    "            start = sql.index(\"from\") + 5\n",
    "            end = sql.find(\" \", start)\n",
    "            return [sql[start:end] if end != -1 else sql[start:]]\n",
    "        return []\n",
    "\n",
    "    def validate_permissions_and_rewrite_sql(self, sql: str, user_id: str) -> dict:\n",
    "        \"\"\"\n",
    "        权限验证和 SQL 重写的整体流程\n",
    "        Args:\n",
    "            sql: 原始 SQL 语句\n",
    "            user_id: 用户ID\n",
    "        Returns:\n",
    "            dict: 包含重写后的 SQL 或错误信息的字典\n",
    "        \"\"\"\n",
    "        print(\"\\n2. Validating permissions and rewriting SQL...\")\n",
    "\n",
    "        # 如果没有提供 user_id，直接返回错误\n",
    "        if not user_id:\n",
    "            return {\n",
    "                'status': 'error',\n",
    "                'message': \"No user_id provided, skipping permission check.\",\n",
    "                'step': 2\n",
    "            }\n",
    "\n",
    "        # 否则进行权限验证和 SQL 重写\n",
    "        result = {'status': 'success'}\n",
    "        result['step'] = 2\n",
    "\n",
    "        try:\n",
    "            rewritten_sql, error = self.rewrite_sql_with_privileges(sql, user_id)\n",
    "            if error:\n",
    "                return {**result, 'status': 'error', 'message': f'SQL权限验证失败: {error}'}\n",
    "            result['rewritten_sql'] = rewritten_sql\n",
    "            print(f\"Rewritten SQL: {rewritten_sql}\")\n",
    "        except Exception as e:\n",
    "            return {**result, 'status': 'error', 'message': f'Error during SQL rewriting: {str(e)}'}\n",
    "\n",
    "        result['message'] = \"Permission check completed successfully.\"\n",
    "        result['sql'] = sql\n",
    "        return result\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestRewriteSQLWithPermissions(unittest.TestCase):\n",
    "    def setUp(self):\n",
    "        self.db_config = {\n",
    "            'user': 'root',  # MySQL 用户名\n",
    "            'password': 'Qqs20001110',  # MySQL 密码\n",
    "            'host': 'localhost',  # MySQL 主机\n",
    "            'database': 'test_db'  # 数据库名称\n",
    "        }\n",
    "        self.instance = SQLRewriter(self.db_config)\n",
    "\n",
    "    def test_validate_permissions_and_rewrite_sql_success(self):\n",
    "        # 使用 self.instance 调用方法\n",
    "        result = self.instance.validate_permissions_and_rewrite_sql(\"SELECT * FROM users WHERE age > 30\", \"user123\")\n",
    "        print(result)\n",
    "\n",
    "        # 断言正确的输出结果\n",
    "        self.assertEqual(result['status'], 'success')\n",
    "        self.assertEqual(result['message'], \"Permission check completed successfully.\")\n",
    "\n",
    "    def test_validate_permissions_and_rewrite_sql_error(self):\n",
    "        # 使用 self.instance 调用方法\n",
    "        result = self.instance.validate_permissions_and_rewrite_sql(\"SELECT * FROM users WHERE age > 30\", \"None\")\n",
    "        print(result)\n",
    "\n",
    "        # 断言错误的输出结果\n",
    "        self.assertEqual(result['status'], 'error')  # 期望错误状态\n",
    "        self.assertIn(\"SQL权限验证失败\", result['message'])  # 期望错误消息包含 SQL 权限验证失败"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "..\n",
      "----------------------------------------------------------------------\n",
      "Ran 2 tests in 0.073s\n",
      "\n",
      "OK\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "2. Validating permissions and rewriting SQL...\n",
      "{'status': 'error', 'step': 2, 'message': 'SQL权限验证失败: 权限不足，无法访问表 users'}\n",
      "\n",
      "2. Validating permissions and rewriting SQL...\n",
      "Rewritten SQL: SELECT name, age FROM users WHERE age > 30\n",
      "{'status': 'success', 'step': 2, 'rewritten_sql': 'SELECT name, age FROM users WHERE age > 30', 'message': 'Permission check completed successfully.', 'sql': 'SELECT * FROM users WHERE age > 30'}\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    unittest.main(argv=['first-arg-is-ignored'], exit=False)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "test-env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
